package com.lucifer.hawkeye.ai.api;

import com.lucifer.hawkeye.ai.common.response.RestResult;
import com.lucifer.hawkeye.ai.service.ChatApiService;
import com.lucifer.hawkeye.ai.service.Nl2SqlService;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestMethod;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import reactor.core.publisher.Flux;
import reactor.core.scheduler.Schedulers;

/**
 * @author lucifer
 * @date 2025/3/29 11:12
 */
@RestController
@RequestMapping("/chat")
public class ChatApi {


    @Resource
    private ChatApiService chatApiService;


    @Resource
    private Nl2SqlService nl2SqlService;


    @RequestMapping(method = {RequestMethod.GET, RequestMethod.POST}, value = "/conversation")
    public RestResult conversation(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "chatType") String chatType,
            @RequestParam(value = "message") String message) {
        return chatApiService.conversation(conversationId, chatType, message);
    }


    @RequestMapping(method = {RequestMethod.GET, RequestMethod.POST}, value = "/conversationV1")
    public RestResult conversationV1(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "conversationType") String conversationType,
            @RequestParam(value = "chatType") String chatType,
            @RequestParam(value = "message") String message) {
        return chatApiService.conversationV1(conversationId,conversationType, chatType, message);
    }


    @RequestMapping(method = {RequestMethod.GET, RequestMethod.POST}, value = "/conversationV2")
    public ResponseBodyEmitter conversationV2(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "chatType") String chatType,
            @RequestParam(value = "message") String message,
            HttpServletResponse servletResponse) {
        ResponseBodyEmitter emitter = new ResponseBodyEmitter();
        servletResponse.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE);
        servletResponse.setCharacterEncoding("UTF-8");
        Flux<ChatResponse> stream = chatApiService.conversationV2(conversationId, chatType, message);
        stream.publishOn(Schedulers.boundedElastic())
                .doOnError(emitter::completeWithError)
                .doOnComplete(emitter::complete)
                .subscribe(data -> {
                    String text = data.getResult().getOutput().getText();
                    try {
                        if (StringUtils.isNotEmpty(text)) {
                            emitter.send(text);
                        }
                    } catch (Exception e) {
                        emitter.completeWithError(e);
                    }
                });
        return emitter;
    }

    @PostMapping("/train")
    public RestResult train(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "trainType") String trainType,
            @RequestParam(value = "question",required = false) String question,
            @RequestParam(value = "content") String content,
            @RequestParam(value = "datasource", required = false) String datasource,
            @RequestParam(value = "tableName", required = false) String tableName) {
        return RestResult.buildSuccessResult(nl2SqlService.train(conversationId,trainType,question,content,datasource,tableName));
    }


    @GetMapping("/generateSql")
    public String generateSql(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "content") String content,
            @RequestParam(value = "datasource", required = false) String datasource,
            @RequestParam(value = "tableName", required = false) String tableName) {
        return nl2SqlService.generateSql(conversationId, content, datasource,tableName);
    }
}
